# !/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Author         :  dax
@Version        :  
------------------------------------
@File           :  dataload.py
@Description    :  
@CreateTime     :  2019/12/18 14:31
------------------------------------
@ModifyTime     :  
"""
import matplotlib.pyplot as plb
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

data_trans = transforms.Compose(
    [transforms.Resize(224),transforms.Grayscale(3), transforms.CenterCrop(224), transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

dataset = ImageFolder(root='./train/', transform=data_trans)

loader = DataLoader(dataset, batch_size=2, shuffle=True)

for img, lab in loader:
    print(img.size())
    g = transforms.ToPILImage()(img[1])
    g.show()
